//      vom_builder.h
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#ifndef VARIABLE_ORDER_MODEL_BUILDER_H
#define VARIABLE_ORDER_MODEL_BUILDER_H

#include "stl-helper.h"
//#include "dna_processing.h"
#include "divsufsort.h"

namespace vom
{
  class sastring : public std::string {
  public:
  sastring(const std::string& s) 
    : std::string(s), T(NULL), SA(NULL), result(0) { build_sa(); }
  sastring(const sastring& s) 
    : std::string(s), T(NULL), SA(NULL), result(0) 
      { 
	copy_pointer(s);
      }
  sastring() : std::string(), T(NULL), SA(NULL), result(0) {}
    sastring& operator=(const std::string& s) {
      this->std::string::operator=(s);
      build_sa();
      return *this;
    }
    sastring& operator=(const sastring& s) {
      this->std::string::operator=(s);
      copy_pointer(s);
      return *this;
    }
    ~sastring() {
      if (T) delete T;
      if (SA) delete SA;
    }
    void sasearch(const std::string& p, int& left, int& size) const {
      if (p.empty()) {
	left = 0;
	size = this->length();
	return;
      }
      sauchar_t *P = NULL;
      str2sachar(p, &P);
      size = sa_search(T, this->length(), P, p.length(), SA, this->length(), &left);
      delete P;
    }

  public:
    sauchar_t *T;
    saidx_t *SA;
    saint_t result;
  protected:
    void build_sa() {
      if (T) delete T;
      if (SA) delete SA;
      T = NULL;
      SA = NULL;
      if (this->empty()) return;
      
      int n = this->length();
      str2sachar(*this, &T);
      SA = new saidx_t[n];
      result = divsufsort(T, SA, n);
      //MSG("DONE 2");
    }

    static void str2sachar(const std::string& s, sauchar_t** T) {
      int n = s.length();
      *T = new sauchar_t[n];
      for (int i = 0; i < n; i++) (*T)[i] = s[i];
      //MSG("DONE 1");
    }
		
    void copy_pointer(const sastring& s)
    {
      int n = s.length();
      T = new sauchar_t[n];
      SA = new saidx_t[n];
      for (int i = 0; i < n; i++) {
	T[i] = s.T[i];
	SA[i] = s.SA[i];				
      }
      result = s.result;
    }
  };
}

namespace vom {
  typedef arma::vec FrequencyTable;
  class node {
  public:
    char c_;
    const char* alphabet_;
    std::vector<FrequencyTable> f_;
    std::vector<node*> child_;
    node* parent_;
    int n_occurence_;
    double lambda_; // weight for interpolated models
  public:
  node(char c, const char* alphabet, node* parent, int periodic) 
    : c_(c), alphabet_(alphabet), 
      f_(periodic, arma::zeros<arma::vec>(strlen(alphabet))), 
      child_(strlen(alphabet), NULL), parent_(parent), n_occurence_(0) {}
    ~node() { BOOST_FOREACH(node* child, child_) delete child; }
    void set_frequency(const vom::sastring& train);
    void set_frequency(const std::vector<vom::sastring>& train_set);
    void set_rand_frequency();
    void normalize();
    std::string str() const;
    int periodic() const { return sz(f_); }
    bool has_table(int pr) const { return sz(f_[pr]) == (int) strlen(alphabet_); }
  };

  class VOM {
  public:
    std::string alphabet;
    node* root_;
    double eta_;
    int count_;
    int periodic_;
    int depth_;

    vi_t count_p_, depth_p_;
		
    VOM(const std::string& alpha, int periodic) 
      : alphabet(alpha), root_(NULL), eta_(0.1), periodic_(periodic),
      count_p_(periodic), depth_p_(periodic) {}
    ~VOM() { delete root_; }

    void build_fom(const vom::sastring& train, int order);
    void build_fom(const std::vector<vom::sastring>& train_set, int order);
    void build_rand_fom(int order);

    int pos(char c);
    char ch(int pos);
		
    double log_prob(const std::string& s);
		
    void print_node(std::ostream& os) { print_node(os, root_, 0); }
    void load_node(std::istream& is) { count_ = 0; load_node(is, &root_, NULL, 0); }
		
    void grow(const std::vector<vom::sastring>& train_set, double eta) 
    { eta_ = eta; count_ = 0; grow_node(root_, train_set); }
    void prune(const std::vector<vom::sastring>& train_set, double eta) 
    { eta_ = eta; prune_node(&root_, train_set); }

    void simulate(int len, std::string& s);
    
    void depth() { 
      count_ = 0; depth_ = 0;
      std::fill(all(count_p_), 0);
      std::fill(all(depth_p_), 0);
      depth_node(root_, 0); 
    }
  protected:
    void build_fom_node(node* N, const vom::sastring& train, int order);
    void build_fom_node(node* N, const std::vector<vom::sastring>& train_set, int order);
    void build_rand_fom_node(node* N, int order);

    void print_node(std::ostream& os, node* N, int depth);
    void load_node(std::istream& is, node** N, node* parent, int depth);

    double log_prob(const std::string& s, int pos);
    const FrequencyTable& f_table(const std::string& s, int pos);
		
    void grow_node(node* N, const std::vector<vom::sastring>& train_set);
    void try_grow(node* N, int pos, const std::vector<vom::sastring>& train_set);
    void prune_node(node** N, const std::vector<vom::sastring>& train_set);
    
    void depth_node(node* N, int depth);
  };
}


#endif
